module aru_ub_wrgen_addr_calc (
    input logic                            clk,
    input logic                            rst_n,
          aru_ub_wrgen_cfg_if.addr_calc_in u_aru_cfg_if,
          aru_payload_if.in                u_aru_pld_if,
          aru_idx_if.in                    u_aru_idx_if,
          aru_sdb_if.in                    u_aru_sdb_if,
          aru_ub_wr_req_if.out             u_ub_wr_req_if
);

    logic cfg_vld;
    logic lst_req_in_instr;
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            cfg_vld <= 'd0;
        end else if (cfg_vld) begin
            if (u_ub_wr_req_if.vld && u_ub_wr_req_if.rdy && lst_req_in_instr) begin
                cfg_vld <= 'd0;
            end
        end else begin
            if (u_aru_cfg_if.vld) begin
                cfg_vld <= 1'b1;
            end
        end
    end
    assign u_aru_cfg_if.rdy = !cfg_vld;
    assign lst_req_in_instr = u_aru_sdb_if.pld.eom && u_aru_sdb_if.pld.eon;

    idx_t slice_m1;
    idx_t slice_n1;
    logic msk      [`M0-1:0][`N0-1:0];

    assign slice_m1 = (u_aru_cfg_if.slice_m + `M0 - 1) / `M0;
    assign slice_n1 = (u_aru_cfg_if.slice_n + `N0 - 1) / `N0;
    always_comb begin
        integer i, j;
        for (i = 0; i < `M0; i = i + 1) begin
            for (j = 0; j < `N0; j = j + 1) begin
                if (i < u_aru_sdb_if.pld.vld_m && j < u_aru_sdb_if.pld.vld_n) begin
                    msk[i][j] = 1'b1;
                end else begin
                    msk[i][j] = 1'b0;
                end
            end
        end
    end

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            u_ub_wr_req_if.ub_addr <= 'd0;
            for (int i = 0; i < `M0; i++) begin
                for (int j = 0; j < `N0; j++) begin
                    u_ub_wr_req_if.dat[i][j] <= '0;
                    u_ub_wr_req_if.msk[i][j] <= '0;
                end
            end
        end else begin
            u_ub_wr_req_if.atomic_mode <= u_aru_cfg_if.atomic_mode;
            u_ub_wr_req_if.ub_addr <= u_aru_cfg_if.ub_addr
                                    + (u_aru_idx_if.pld.n1_idx * slice_m1 * `M0 + (u_aru_idx_if.pld.m1_idx * `M0 + u_aru_idx_if.pld.p_idx * `P_ARU)) * `N0 * 2;
            // 将ARU输出数据重新排列到UB写入格式
            for (int i = 0; i < `M0; i++) begin
                for (int j = 0; j < `N0; j++) begin
                    // 如果是`P_ARU个处理单元,需要适当映射数据
                    if (i < `P_ARU) begin
                        u_ub_wr_req_if.dat[i][j] <= u_aru_pld_if.dat[i*`N0+j];
                    end else begin
                        u_ub_wr_req_if.dat[i][j] <= '0;
                    end
                    u_ub_wr_req_if.msk[i][j] <= msk[i][j];
                end
            end
        end
    end

    logic has_data;
    logic up_vld, up_rdy, down_vld, down_rdy;
    assign down_rdy = u_ub_wr_req_if.rdy;
    assign up_vld   = u_aru_idx_if.vld && u_aru_sdb_if.vld && u_aru_pld_if.vld;
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            has_data <= 'd0;
        end else if (has_data) begin
            if (~up_vld && down_rdy) begin
                has_data <= 'd0;
            end
        end else begin
            if (up_vld) begin
                has_data <= 'd1;
            end
        end
    end
    assign down_vld           = has_data;
    assign up_rdy             = (~has_data) || down_rdy;

    assign u_ub_wr_req_if.vld = cfg_vld && down_vld;
    assign u_aru_idx_if.rdy   = up_rdy && cfg_vld;
    assign u_aru_sdb_if.rdy   = up_rdy && cfg_vld;
    assign u_aru_pld_if.rdy   = up_rdy && cfg_vld;

endmodule
